#include "Graphics/PiccoloVKSwapchain.h"

#include <Graphics/PiccoloVKDevice.h>
#include <Graphics/PiccoloVKGraphicsContext.h>
#include <Graphics/PiccoloVKQueue.h>

namespace Piccolo
{
    PiccoloVKSwapchain::PiccoloVKSwapchain(PiccoloVKGraphicsContext *context, PiccoloVKDevice *device) : mContext(context), mDevice(device)
    {
        RecreateSwapchain();
    }

    PiccoloVKSwapchain::~PiccoloVKSwapchain()
    {
        VK_D(SwapchainKHR, mDevice->GetLogicDevice(), mSwapchain);
    }

    bool PiccoloVKSwapchain::RecreateSwapchain()
    {
        LOG_D("-------------------------");
        SetupSurfaceCapabilities();
        LOG_D("currentExtent : {0} x {1}", mSurfaceInfo.capabilities.currentExtent.width, mSurfaceInfo.capabilities.currentExtent.height);
        LOG_D("surfaceFormat : {0}", vk_format_string(mSurfaceInfo.surfaceFormat.format));
        LOG_D("presentMode   : {0}", vk_present_mode_string(mSurfaceInfo.presentMode));
        LOG_D("-------------------------");

        uint32_t imageCount = mDevice->GetSettings().swapchainImageCount;
        if (imageCount < mSurfaceInfo.capabilities.minImageCount && mSurfaceInfo.capabilities.minImageCount > 0)
        {
            imageCount = mSurfaceInfo.capabilities.minImageCount;
        }
        if (imageCount > mSurfaceInfo.capabilities.maxImageCount && mSurfaceInfo.capabilities.maxImageCount > 0)
        {
            imageCount = mSurfaceInfo.capabilities.maxImageCount;
        }

        VkSharingMode imageSharingMode;
        uint32_t queueFamilyIndexCount;
        uint32_t pQueueFamilyIndices[2] = {0, 0};
        if (mContext->IsSameGraphicPresentQueueFamily())
        {
            imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
            queueFamilyIndexCount = 0;
        }
        else
        {
            imageSharingMode = VK_SHARING_MODE_CONCURRENT;
            queueFamilyIndexCount = 2;
            pQueueFamilyIndices[0] = mContext->GetGraphicQueueFamilyInfo().queueFamilyIndex;
            pQueueFamilyIndices[1] = mContext->GetPresentQueueFamilyInfo().queueFamilyIndex;
        }

        VkSwapchainKHR oldSwapchain = mSwapchain;

        VkSwapchainCreateInfoKHR swapchainInfo =
        {
            .sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
            .pNext = nullptr,
            .flags = 0,
            .surface = mContext->GetSurface(),
            .minImageCount = imageCount,
            .imageFormat = mSurfaceInfo.surfaceFormat.format,
            .imageColorSpace = mSurfaceInfo.surfaceFormat.colorSpace,
            .imageExtent = mSurfaceInfo.capabilities.currentExtent,
            .imageArrayLayers = 1,
            .imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
            .imageSharingMode = imageSharingMode,
            .queueFamilyIndexCount = queueFamilyIndexCount,
            .pQueueFamilyIndices = pQueueFamilyIndices,
            .preTransform = VK_SURFACE_TRANSFORM_IDENTITY_BIT_KHR,
            .compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
            .presentMode = mSurfaceInfo.presentMode,
            .clipped = VK_FALSE,
            .oldSwapchain = oldSwapchain
        };
        VkResult ret = vkCreateSwapchainKHR(mDevice->GetLogicDevice(), &swapchainInfo, nullptr, &mSwapchain);
        if (ret != VK_SUCCESS)
        {
            LOG_E("{0} : {1}", __FUNCTION__, vk_result_string(ret));
            return false;
        }
        LOG_D("Swapchain {0} : old: {1}, new: {2}, image count: {3}, format: {4}, present mode : {5}", __FUNCTION__, (void*)oldSwapchain, (void*)mSwapchain, imageCount,
            vk_format_string(mSurfaceInfo.surfaceFormat.format), vk_present_mode_string(mSurfaceInfo.presentMode));

        uint32_t swapchainImageCount;
        vkGetSwapchainImagesKHR(mDevice->GetLogicDevice(), mSwapchain, &swapchainImageCount, nullptr);
        mSwapchainImages.resize(swapchainImageCount);
        ret = vkGetSwapchainImagesKHR(mDevice->GetLogicDevice(), mSwapchain, &swapchainImageCount, mSwapchainImages.data());

        return ret == VK_SUCCESS;
    }

    VkResult PiccoloVKSwapchain::AcquireImage(int32_t* outImageIndex, VkSemaphore semaphore, VkFence fence)
    {
        uint32_t imageIndex;
        VkResult ret = vkAcquireNextImageKHR(mDevice->GetLogicDevice(), mSwapchain, UINT64_MAX, semaphore, fence, &imageIndex);
        if (fence != VK_NULL_HANDLE)
        {
            CALL_VK(vkWaitForFences(mDevice->GetLogicDevice(), 1, &fence, VK_FALSE, UINT64_MAX));
            CALL_VK(vkResetFences(mDevice->GetLogicDevice(), 1, &fence));
        }
        if (ret == VK_SUCCESS || ret == VK_SUBOPTIMAL_KHR)
        {
            *outImageIndex = imageIndex;
            mCurrentImageIndex = imageIndex;
        }
        return ret;
    }

    VkResult PiccoloVKSwapchain::Present(int32_t imageIndex, const std::vector<VkSemaphore>& waitSemaphores)
    {
        VkPresentInfoKHR presentInfo =
        {
            .sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR,
            .pNext = nullptr,
            .waitSemaphoreCount = static_cast<uint32_t>(waitSemaphores.size()),
            .pWaitSemaphores = waitSemaphores.data(),
            .swapchainCount = 1,
            .pSwapchains = &mSwapchain,
            .pImageIndices = reinterpret_cast<const uint32_t*>(&imageIndex)
        };
        VkResult ret = vkQueuePresentKHR(mDevice->GetFirstPresentQueue()->GetQueue(), &presentInfo);
        mDevice->GetFirstPresentQueue()->WaitIdle();
        return ret;
    }

    /**
     * 设置surface的一些参数
     */
    void PiccoloVKSwapchain::SetupSurfaceCapabilities()
    {
        // capabilities
        vkGetPhysicalDeviceSurfaceCapabilitiesKHR(mContext->GetPhysicalDevice(), mContext->GetSurface(), &mSurfaceInfo.capabilities);

        PiccoloVKSettings settings = mDevice->GetSettings();

        // format
        uint32_t formatCount;
        CALL_VK(vkGetPhysicalDeviceSurfaceFormatsKHR(mContext->GetPhysicalDevice(), mContext->GetSurface(), &formatCount, nullptr));
        if (formatCount == 0)
        {
            LOG_E("{0} : num of surface format is 0", __FUNCTION__);
            return;
        }
        VkSurfaceFormatKHR formats[formatCount];
        CALL_VK(vkGetPhysicalDeviceSurfaceFormatsKHR(mContext->GetPhysicalDevice(), mContext->GetSurface(), &formatCount, formats));

        // 查找是否有期望的格式
        uint32_t foundFormatIndex = -1;
        for (int i = 0; i < formatCount; ++i)
        {
            if (formats[i].format == settings.surfaceFormat && formats[i].colorSpace == VK_COLORSPACE_SRGB_NONLINEAR_KHR)
            {
                foundFormatIndex = i;
                break;
            }
        }
        if (foundFormatIndex == -1)
        {
            foundFormatIndex = 0;
        }
        mSurfaceInfo.surfaceFormat = formats[foundFormatIndex];

        // present mode
        uint32_t presentModeCount;
        CALL_VK(vkGetPhysicalDeviceSurfacePresentModesKHR(mContext->GetPhysicalDevice(), mContext->GetSurface(), &presentModeCount, nullptr));
        if (presentModeCount == 0)
        {
            LOG_E("{0} : num of surface present mode is 0", __FUNCTION__);
            return;
        }
        VkPresentModeKHR presentModes[presentModeCount];
        CALL_VK(vkGetPhysicalDeviceSurfacePresentModesKHR(mContext->GetPhysicalDevice(), mContext->GetSurface(), &presentModeCount, presentModes));
        // 期望的呈现模式
        VkPresentModeKHR preferedPresentMode = mDevice->GetSettings().presentMode;
        uint32_t foundPresentModeIndex = -1;
        for (uint32_t i = 0; i < presentModeCount; ++i)
        {
            if (presentModes[i] == preferedPresentMode)
            {
                foundPresentModeIndex = i;
                break;
            }
        }
        if (foundPresentModeIndex > 0)
        {
            mSurfaceInfo.presentMode = presentModes[foundPresentModeIndex];
        }
        else
        {
            mSurfaceInfo.presentMode = presentModes[0];
        }
    }
}
